﻿using System;
using System.Management;
using System.Diagnostics;
using System.Collections.Generic;

namespace HyperV
{
    public partial class VMHelper
    {
        private static ManagementObject GetLastVirtualSystemSnapshot(ManagementObject vm)
        {
            ManagementObjectCollection settings = vm.GetRelated(
                "Msvm_VirtualSystemsettingData",
                "Msvm_PreviousSettingData",
                null,
                null,
                "SettingData",
                "ManagedElement",
                false,
                null);

            ManagementObject virtualSystemsetting = null;
            foreach (ManagementObject setting in settings)
            {
                Debug.WriteLine(setting.Path.Path);
                Debug.WriteLine(setting["ElementName"]);
                virtualSystemsetting = setting;
            }

            return virtualSystemsetting;
        }

        private static ManagementBaseObject ApplyVirtualSystemSnapshot(ManagementScope scope, string vmName)
        {
            ManagementObject virtualSystemService = Utility.GetServiceObject(scope, "Msvm_VirtualSystemManagementService");

            ManagementBaseObject inParams = virtualSystemService.GetMethodParameters("ApplyVirtualSystemSnapshot");

            ManagementObject vm = Utility.GetTargetComputer(vmName, scope);

            ManagementObject vmSnapshot = GetLastVirtualSystemSnapshot(vm);

            inParams["SnapshotSettingData"] = vmSnapshot.Path.Path;

            inParams["ComputerSystem"] = vm.Path.Path;

            ManagementBaseObject outParams = virtualSystemService.InvokeMethod("ApplyVirtualSystemSnapshot", inParams, null);

            return outParams;
        }

        public void Revert(params String[] vmNames) {
            Revert(Scope, vmNames);
        }

        public static void Revert(ManagementScope scope, params String[] vmNames)
        {
            Dictionary<String, ManagementBaseObject> results = new Dictionary<String, ManagementBaseObject>();

            foreach (String vmName in vmNames)
            {
                Stop(scope, vmName);

                ManagementBaseObject outParams = ApplyVirtualSystemSnapshot(scope, vmName);
                if ((UInt32)outParams["ReturnValue"]==ReturnCode.Completed)
                    Debug.WriteLine(vmName + ": Snapshot was applied successfully.");
                else if ((UInt32)outParams["ReturnValue"]==ReturnCode.Started)
                    results.Add(vmName, outParams);
                else Debug.WriteLine(vmName + ": Failed to apply snapshot. ErrorCode: " + outParams["ReturnValue"]);
            }

            foreach (KeyValuePair<String, ManagementBaseObject> result in results)
            {

                if (Utility.JobCompleted(result.Value, scope))
                {
                    Debug.WriteLine(result.Key + ": Snapshot was applied successfully.");
                }
                else
                {
                    Debug.WriteLine(result.Key + ": Failed to apply snapshot.");
                }
            }

            foreach (String vmName in vmNames)
            {
                Start(scope, vmName);
            }
        }
    }
}